# coding:utf-8

import csv
import os
import codecs
import cv2
import numpy as np
import math

from Utils import utils
from Utils.DataLoader import IITDLoader, CASIALoader, TongjiLoader
from Zhang.zhangROI import ZhangROIExtractor

src = np.array([[5.79 / 12.67 * 480, 8.15 / 12.91 * 640], [7.86 / 12.67 * 480, 9.26 / 12.91 * 640],
                [4.47 / 12.67 * 480, 14.37 / 12.91 * 640]],
               np.float32)
dst = np.array([[8.04 / 16.91 * 480, 12.19 / 22.55 * 640], [9.43 / 16.91 * 480, 8.55 / 22.55 * 640],
                [16.26 / 16.91 * 480, 14.56 / 22.55 * 640]],
               np.float32)
M = cv2.getAffineTransform(src, dst)


def transform_kps(x0, x2, y0, y2):
    # cv2.warpAffine()
    # x towards down; y towards right
    ori = np.array([[x0, y0], [x2, y2]], np.float32)
    W = M[:, :-1]
    b = M[:, -1] + (65, 0)
    res = np.matmul(W, ori.transpose())
    x_new0, y_new0 = res[:, 0] + b
    x_new2, y_new2 = res[:, 1] + b

    # x0 = x0 * 10 / 3
    # x2 = x2 * 10 / 3
    # y0 = 2.45625 * y0
    # y2 = 2.45625 * y2
    #
    # y0 = y0 - 186
    # y2 = y2 - 186
    #
    # y_new0 = 1600 - x0
    # y_new2 = 1600 - x2
    # x_new0 = y0
    # x_new2 = y2
    #
    # x_new0 = x_new0 * 0.4
    # x_new2 = x_new2 * 0.4
    # y_new0 = y_new0 * 0.4
    # y_new2 = y_new2 * 0.4

    return int(x_new0), int(x_new2), int(y_new0), int(y_new2)


def extractClear_IITD(img_ori, (x0, x2, y0, y2), file_name, path_save, rio_extractor):
    x0, x2, y0, y2 = transform_kps(x0, x2, y0, y2)
    img = cv2.resize(img_ori, dsize=(640, 480))
    img_sharp = img
    # kernel_size = 5.0
    # cv2.imwrite(os.path.join(path_save, "sharpen" + str(-1) + "_" + file_name), img_sharp)
    # for i in range(8):
    #     g_kernel = cv2.getGaborKernel(ksize=(int(kernel_size), int(kernel_size)), sigma=kernel_size / 3,
    #                                   theta=i * 2 * math.pi / 8, lambd=kernel_size / 3 / 20,
    #                                   gamma=1, psi=0)
    #     g_kernel = g_kernel / kernel_size / kernel_size
    #     img_sharp = cv2.filter2D(src=img_sharp, kernel=g_kernel, ddepth=cv2.CV_64F)
    #     cv2.imwrite(os.path.join(path_save, "sharpen" + str(i) + "_" + file_name), img_sharp)
    img_sharp = np.uint8(img_sharp)
    roi = rio_extractor.extract_kp(img_sharp, x0, x2, y0, y2)
    roisum = roi.sum()
    if roisum > 100:
        cv2.imwrite(os.path.join(path_save, file_name), roi)
    else:
        print "bad ROI!"


def extractClearNNROI(dataLoader, ROI_output_dir, roie):
    dataLoader.iterate(extractClear_IITD, ROI_output_dir, roie)


if __name__ == '__main__':
    opt = dict({
        # "NN_input_dir": '/home/yjy/PycharmProjects/HRNet-Facial-Landmark-Detection/data/iitd/images',
        # "predicted_labels": '/home/yjy/PycharmProjects/HRNet-Facial-Landmark-Detection/output/IITD/' \
        #                     'palm_alignment_iitd_hrnet_w18/predictions_all.csv',
        "predicted_labels": '/home/yjy/dataset/iitd_predictions_all.csv',
        "Clear_images_input_dir": "/home/yjy/dataset/IITD/IITD Palmprint V1/Left Hand",
        "Clear_ROI_output_dir": "/home/yjy/dataset/IITD-ROI-NN-clear",
    })

    roie = ZhangROIExtractor()
    dataLoader = IITDLoader(opt["Clear_images_input_dir"], opt["predicted_labels"])
    extractClearNNROI(dataLoader, opt["Clear_ROI_output_dir"], roie)
